//
//  FuOcrLite.m
//  FunnyReader
//
//  Created by yht on 2025/1/21.
//

#include "FuOcrLite.h"
#include "AngleNet.h"
#include "CrnnNet.h"
#include "DbNet.h"
#include "OcrUtils.h"
#import "ImageUtils.h"

@implementation FunOcrResult
@end


@interface FuOcrLite()
{
    AngleNet angleNet;
    CrnnNet crnnNet;
    DbNet dbNet;
}
//@property(nonatomic, strong)AngleNet *angleNet;


@end

@implementation FuOcrLite

- (instancetype)init
{
    self = [super init];
    if (self) {
        [self initialize:1];
    }
    return self;
}


- (void)initialize:(int)numOfThread {

    FLogger("--- Init AngleNet ---\n");
    
    BOOL retAngleNet = [self initAngleNetModel:numOfThread];

    FLogger("--- Init CrnnNet ---\n");
    BOOL retCrnnNet = [self initCrnnNetModel:numOfThread];
    
    FLogger("--- Init DbNet ---\n");
    BOOL retDbNet = [self initDbNetModel:numOfThread];

    if (!retDbNet || !retAngleNet || !retCrnnNet) {
        FLogger("Init Models Failed! %d  %d  %d", retDbNet, retAngleNet, retCrnnNet);
    }

    FLogger("Init Models Success!");
    
}

- (BOOL)initAngleNetModel:(int)numOfThread {
    angleNet.setNumThread(numOfThread);
    NSString *binPath = [self loadModelPath:@"angle_op.bin"];
    NSString *paramPath = [self loadModelPath:@"angle_op.param"];
    return angleNet.initModel(paramPath.UTF8String, binPath.UTF8String);
}

- (BOOL)initCrnnNetModel:(int)numOfThread {
    crnnNet.setNumThread(numOfThread);
    NSString *binPath = [self loadModelPath:@"crnn_lite_op.bin"];
    NSString *paramPath = [self loadModelPath:@"crnn_lite_op.param"];
    NSString *keyPath = [self loadModelPath:@"keys.txt"];
    return crnnNet.initModel(paramPath.UTF8String, binPath.UTF8String, keyPath.UTF8String);
}

- (BOOL)initDbNetModel:(int)numOfThread {
    dbNet.setNumThread(numOfThread);
    NSString *binPath = [self loadModelPath:@"dbnet_op.bin"];
    NSString *paramPath = [self loadModelPath:@"dbnet_op.param"];
    return dbNet.initModel(paramPath.UTF8String, binPath.UTF8String);
}


- (NSString*)loadModelPath:(NSString*)model {
    NSBundle *mainBundle = [NSBundle bundleForClass:[self class]];
    NSBundle *modelBundle = [NSBundle bundleWithURL:[mainBundle URLForResource:@"Resources" withExtension:@"bundle"]];
    NSString *modelPath = [modelBundle pathForResource:model ofType:nil];
    return modelPath;
}


- (OcrResult)detect:(cv::Mat&)src originRect:(cv::Rect&)originRect scale:(ScaleParam&)scale boxScoreThresh:(float)boxScoreThresh boxThresh:(float)boxThresh unClipRatio:(float)unClipRatio doAngle:(bool)doAngle mostAngle:(bool)mostAngle {
    cv::Mat textBoxPaddingImg = src.clone();
    int thickness = getThickness(src);

    FLogger("=====Start detect=====");
    FLogger("ScaleParam(sw:%d,sh:%d,dw:%d,dh:%d,%f,%f)", scale.srcWidth, scale.srcHeight,
           scale.dstWidth, scale.dstHeight,
           scale.ratioWidth, scale.ratioHeight);

    FLogger("---------- step: dbNet getTextBoxes ----------");
    double startTime = getCurrentTime();
//    UIImage *image = [ImageUtils imageWithCVMat:src];
    std::vector<TextBox> textBoxes = dbNet.getTextBoxes(src, scale, boxScoreThresh, boxThresh, unClipRatio);
    FLogger("TextBoxesSize(%ld)", textBoxes.size());
    if (textBoxes.size() == 0) {
        return OcrResult{};
    }
    double endDbNetTime = getCurrentTime();
    double dbNetTime = endDbNetTime - startTime;
    FLogger("dbNetTime(%fms)", dbNetTime);

    for (int i = 0; i < textBoxes.size(); ++i) {
        FLogger("TextBox[%d][score(%f),[x: %d, y: %d], [x: %d, y: %d], [x: %d, y: %d], [x: %d, y: %d]]",
               i,
               textBoxes[i].score,
               textBoxes[i].boxPoint[0].x, textBoxes[i].boxPoint[0].y,
               textBoxes[i].boxPoint[1].x, textBoxes[i].boxPoint[1].y,
               textBoxes[i].boxPoint[2].x, textBoxes[i].boxPoint[2].y,
               textBoxes[i].boxPoint[3].x, textBoxes[i].boxPoint[3].y);
    }

    FLogger("---------- step: drawTextBoxes ----------");
//    drawTextBoxes(textBoxPaddingImg, textBoxes, thickness);

    //---------- getPartImages ----------
    std::vector<cv::Mat> partImages = [self getPartImages:src textBoxes:textBoxes];
    
    //Rotate partImgs
//    for (int i = 0; i < partImages.size(); ++i) {
//        UIImage *image = [ImageUtils imageWithCVMat:partImages[i]];
//        FLogger("hi, h");
//    }

    FLogger("---------- step: angleNet getAngles ----------");
    std::vector<Angle> angles;
    angles = angleNet.getAngles(partImages, doAngle, mostAngle);
    if (angles.size() == 0) {
        return OcrResult{};
    }

    //Log Angles
    for (int i = 0; i < angles.size(); ++i) {
        FLogger("angle[%d][index(%d), score(%f), time(%fms)]", i, angles[i].index, angles[i].score, angles[i].time);
    }

    //Rotate partImgs
    for (int i = 0; i < partImages.size(); ++i) {
        if (angles[i].index == 0) {
            partImages.at(i) = matRotateClockWise180(partImages[i]);
        }
//        UIImage *image = [ImageUtils imageWithCVMat:partImages[i]];
//        FLogger("hi, h");
    }

    FLogger("---------- step: crnnNet getTextLine ----------");
    std::vector<TextLine> textLines = crnnNet.getTextLines(partImages);
    if (textLines.size() == 0) {
        return OcrResult{};
    }
    //Log TextLines
    for (int i = 0; i < textLines.size(); ++i) {
        FLogger("textLine[%d](%s)", i, textLines[i].text.c_str());
        std::ostringstream txtScores;
        for (int s = 0; s < textLines[i].charScores.size(); ++s) {
            if (s == 0) {
                txtScores << textLines[i].charScores[s];
            } else {
                txtScores << " ," << textLines[i].charScores[s];
            }
        }
        FLogger("textScores[%d]{%s}", i, std::string(txtScores.str()).c_str());
        FLogger("crnnTime[%d](%fms)", i, textLines[i].time);
    }

    std::vector<TextBlock> textBlocks;
    for (int i = 0; i < textLines.size(); ++i) {
        std::vector<cv::Point> boxPoint = std::vector<cv::Point>(4);
        int padding = originRect.x;//padding conversion
        boxPoint[0] = cv::Point(textBoxes[i].boxPoint[0].x - padding, textBoxes[i].boxPoint[0].y - padding);
        boxPoint[1] = cv::Point(textBoxes[i].boxPoint[1].x - padding, textBoxes[i].boxPoint[1].y - padding);
        boxPoint[2] = cv::Point(textBoxes[i].boxPoint[2].x - padding, textBoxes[i].boxPoint[2].y - padding);
        boxPoint[3] = cv::Point(textBoxes[i].boxPoint[3].x - padding, textBoxes[i].boxPoint[3].y - padding);
        TextBlock textBlock{boxPoint, textBoxes[i].score, angles[i].index, angles[i].score,
                            angles[i].time, textLines[i].text, textLines[i].charScores, textLines[i].time,
                            angles[i].time + textLines[i].time};
        textBlocks.emplace_back(textBlock);
    }

    double endTime = getCurrentTime();
    double fullTime = endTime - startTime;
    FLogger("=====End detect=====");
    FLogger("FullDetectTime(%fms)", fullTime);

    //cropped to original size
    cv::Mat textBoxImg;
    if (originRect.x > 0 && originRect.y > 0) {
        textBoxPaddingImg(originRect).copyTo(textBoxImg);
    } else {
        textBoxImg = textBoxPaddingImg;
    }

    std::string strRes;
    for (int i = 0; i < textBlocks.size(); ++i) {
        strRes.append(textBlocks[i].text);
        strRes.append("\n");
    }

    return OcrResult{dbNetTime, textBlocks, textBoxImg, fullTime, strRes};
}

- (std::vector<cv::Mat>) getPartImages:(cv::Mat&)src textBoxes:(std::vector<TextBox>&)textBoxes {
    std::vector<cv::Mat> partImages;
    for (int i = 0; i < textBoxes.size(); ++i) {
        cv::Mat partImg = getRotateCropImage(src, textBoxes[i].boxPoint);
        partImages.emplace_back(partImg);
    }
    return partImages;
}

- (void)ocrDetection:(UIImage*)input padding:(int)padding maxSideLen:(int)maxSideLen boxScoreThresh:(float)boxScoreThresh boxThresh:(float)boxThresh unClipRatio:(float)unClipRatio doAngle:(BOOL)doAngle mostAngle:(BOOL)mostAngle {

    FLogger("padding(%d),maxSideLen(%d),boxScoreThresh(%f),boxThresh(%f),unClipRatio(%f),doAngle(%d),mostAngle(%d)",
           padding, maxSideLen, boxScoreThresh, boxThresh, unClipRatio, doAngle, mostAngle);
//    input = [ImageUtils fixOrientationOfImage:input];
    cv::Mat imgBGR, imgOut;
    imgBGR = [ImageUtils UIImageToMat:input];
    if (imgBGR.empty()) {
        return;
    }
//    UIImage *image = [ImageUtils imageWithCVMat:imgBGR];
    int originMaxSide = (std::max)(imgBGR.cols, imgBGR.rows);
    int resize;
    if (maxSideLen <= 0 || maxSideLen > originMaxSide) {
        resize = originMaxSide;
    } else {
        resize = maxSideLen;
    }
    resize += 2*padding;
    cv::Rect paddingRect(padding, padding, imgBGR.cols, imgBGR.rows);
    cv::Mat paddingSrc = [self makePadding:imgBGR padding:padding];
//    UIImage *image = [ImageUtils imageWithCVMat:paddingSrc];
    //按比例缩小图像，减少文字分割时间
//    ScaleParam s = getScaleParam(paddingSrc, resize);//例：按长或宽缩放 src.cols=不缩放，src.cols/2=长度缩小一半
    ScaleParam s = getScaleParam(paddingSrc, 600);
                          
                          
    OcrResult ocrResult = [self detect:paddingSrc originRect:paddingRect scale:s boxScoreThresh:boxScoreThresh boxThresh:boxThresh
                           unClipRatio:unClipRatio doAngle:doAngle mostAngle:mostAngle];

    if (self.delegate!=nil) {
        FunOcrResult *result = [self convertResult:ocrResult];
        [self.delegate doneDetection:result];
    }

}

- (FunOcrResult*)convertResult:(OcrResult&)result {
    FunOcrResult *nr = [FunOcrResult new];
    nr.text = [NSString stringWithUTF8String:result.strRes.c_str()];
    return nr;
}

- (cv::Mat)makePadding:(cv::Mat&)src padding:(const int)padding {
    if (padding <= 0) return src;
    cv::Scalar paddingScalar = {255, 255, 255};
    cv::Mat paddingSrc;
    cv::copyMakeBorder(src, paddingSrc, padding, padding, padding, padding, cv::BORDER_ISOLATED,
                       paddingScalar);
    return paddingSrc;
}


@end
